#include <stdint.h>
#include <string.h>
#include <stdio.h>
#include "memcheck.h"
#include "memcheck/common.h"

static inline void emit8(uint8_t **buf, uint8_t v) { **buf = v; *buf += 1; }
static inline void emit64le(uint8_t **buf, uint64_t v) { *(uint64_t*)(*buf) = v; *buf += 8; }

static size_t decode_x86_64(const uint8_t *p, int *is_mem, int *is_store, size_t *size)
{
    const uint8_t *s = p;
    int rex_w = 0, has_66 = 0;
    *is_mem = 0; *is_store = 0; *size = 0;
    while (1) {
        uint8_t b = *s;
        if (b == 0x66) { has_66 = 1; s++; continue; }
        if (b == 0xF0 || b == 0xF2 || b == 0xF3 ||
            b == 0x2E || b == 0x36 || b == 0x3E || b == 0x26 ||
            b == 0x64 || b == 0x65 || b == 0x67) { s++; continue; }
        if (b >= 0x40 && b <= 0x4F) { if (b & 0x8) rex_w = 1; s++; continue; }
        break;
    }
    uint8_t op = *s++;
    if (op == 0xC3) return (size_t)(s - p);
    if (op == 0xE8 || op == 0xE9) { s += 4; return (size_t)(s - p); }
    if (op == 0xEB) { s += 1; return (size_t)(s - p); }
    if (op >= 0x70 && op <= 0x7F) { s += 1; return (size_t)(s - p); }
    if (op == 0x0F) { uint8_t op2 = *s++; if (op2 >= 0x80 && op2 <= 0x8F) { s += 4; return (size_t)(s - p); } }
    uint8_t modrm = *s++;
    uint8_t mod = (modrm >> 6) & 3;
    uint8_t rm  = (modrm & 7);
    if (rm == 4 && mod != 3) { uint8_t sib = *s++; uint8_t base = sib & 7; if (mod == 0 && base == 5) s += 4; }
    if (mod == 1) s += 1; else if (mod == 2) s += 4; else if (mod == 0 && rm == 5) s += 4;
    if (op == 0x8B || op == 0x89) { if (mod != 3) { *is_mem = 1; *is_store = (op == 0x89); if (rex_w) *size = 8; else if (has_66) *size = 2; else *size = 4; } }
    return (size_t)(s - p);
}

static void emit_mov_imm64_reg(uint8_t **out, int reg, uint64_t imm)
{
    emit8(out, 0x48 | ((reg >> 3) & 1));
    emit8(out, 0xB8 + (reg & 7));
    emit64le(out, imm);
}

static void emit_call_abs(uint8_t **out, uint64_t target)
{
    emit8(out, 0x49); emit8(out, 0xBB); emit64le(out, target); // mov r11, imm64
    emit8(out, 0x41); emit8(out, 0xFF); emit8(out, 0xD3);     // call r11
}

static void emit_push_caller_saved(uint8_t **out)
{
    emit8(out, 0x9C);
    emit8(out, 0x50); emit8(out, 0x51); emit8(out, 0x52);
    emit8(out, 0x56); emit8(out, 0x57);
    emit8(out, 0x41); emit8(out, 0x50);
    emit8(out, 0x41); emit8(out, 0x51);
    emit8(out, 0x41); emit8(out, 0x52);
    emit8(out, 0x41); emit8(out, 0x53);
}

static void emit_pop_caller_saved(uint8_t **out)
{
    emit8(out, 0x41); emit8(out, 0x5B);
    emit8(out, 0x41); emit8(out, 0x5A);
    emit8(out, 0x41); emit8(out, 0x59);
    emit8(out, 0x41); emit8(out, 0x58);
    emit8(out, 0x5F); emit8(out, 0x5E);
    emit8(out, 0x5A); emit8(out, 0x59); emit8(out, 0x58);
    emit8(out, 0x9D);
}

void dbi_log_mem_access(uint64_t pc, uint64_t ea, int is_store, size_t size)
{
    dprintf(2, "[memcheck:x64] rip=0x%lx ea=0x%lx %s size=%zu\n",
            (unsigned long)pc, (unsigned long)ea, is_store?"STORE":"LOAD", size);
}

static void emit_call_log(uint8_t **out, uint64_t pc, uint64_t ea, int is_store, size_t size)
{
    emit_push_caller_saved(out);
    // SysV ABI: rdi, rsi, rdx, rcx
    emit_mov_imm64_reg(out, 7, pc);
    emit_mov_imm64_reg(out, 6, ea);
    emit_mov_imm64_reg(out, 2, (uint64_t)is_store);
    emit_mov_imm64_reg(out, 1, (uint64_t)size);
    emit_call_abs(out, (uint64_t)&dbi_log_mem_access);
    emit_pop_caller_saved(out);
}

void *dbi_translate_function(void *orig_fn)
{
    const uint8_t *in = (const uint8_t *)orig_fn;
    uint8_t *out = (uint8_t *)mc_code_cache_alloc(8192);
    if (!out) return NULL;
    uint8_t *cursor = out;
    for (int i = 0; i < 4096; i++) {
        int is_mem = 0, is_store = 0; size_t sz = 0;
        size_t len = decode_x86_64(in, &is_mem, &is_store, &sz);
        if (len == 0) break;
        if (is_mem) emit_call_log(&cursor, (uint64_t)in, 0, is_store, sz);
        memcpy(cursor, in, len);
        cursor += len;
        if (in[0] == 0xC3) break;
        in += len;
    }
    mc_clear_icache(out, cursor);
    return out;
}


